import pylab as plt
import sys
import os
import math
from params import *
import scipy
import scipy.stats

def scatter(x,y,x_lbl,y_lbl,name):
    plt.figure()
    plt.scatter(x,y)
    plt.grid(True)
    plt.xlabel(x_lbl)
    plt.ylabel(y_lbl)
    plt.savefig(name)
    plt.close()

def bb(data1,data2,m1,m2,log=False):
    print min(data1), max(data1)
    print m1,m2
    if log:
        m1,m2 = map(math.log10,(m1,m2))
    values = {}
    bins = 20
    factor = float(bins)/(m2-m1)
    for x,y in zip(data1,data2):
        if log:
            x = math.log10(x)
        if not m1 <= x <= m2:
            continue
        index = int(x*factor)
        index = float(index)/factor
        if log:
            index = 10**index
        values.setdefault(index,[]).append(y)

    x,y = [],[]
    s = []
    for key in sorted(values):
        x.append(key)
        d = values[key]
        avg = 0.0
        std = 0.0
        if len(d) > 0:
            avg = float(sum(d))/len(d) 
            var = float(sum((i-avg)**2 for i in d))/len(d)
            std = var**0.5
        y.append(avg)
        s.append(std)
    return x,y,s

DIR = os.path.join(WORKDIR, 'gowalla', 'results', 'entropy')
plots = {}
for file in os.listdir(DIR):
    name,ext = os.path.splitext(file)
    if not name.startswith('place_diversity_correlation'):
        continue
    date = name.split('_')[-1]
    month = int(date.split('-')[1])
    print name, date, month

    x1,x2,x3  = [],[],[]
    y1,y2 = [],[]
    for line in open(os.path.join(DIR,file)):
        try:
            t = map(float,line.strip().split())
        except:
            pass
        k,cat,num_users,num_checkins,prob,d,e = t
        x1.append(num_users)
        x2.append(prob)
        x3.append(num_checkins)
        y1.append(d)
        y2.append(e)
    plots[month] = (x1,x2,y1,y2,x3)

#markers = list('+x.v')
#FIG_AXES2 = [0.20,0.2,0.95-0.20,0.95-0.2]
#plt.figure()
#plt.clf()
#plt.axes(FIG_AXES2)
#lbl = []
#for month,m in zip(sorted(plots),markers):
#    x1,x2,y1,y2,x3 = plots[month]
#    a1,b1,s1 = bb(y1,x2,1.001,10)
#    plt.semilogy(a1,b1,'k%s-'%m)
#    lbl.append('Snapshot %d'%(month-4))
#plt.legend(lbl,loc='lower left')
#plt.grid(True)
#plt.xlabel('Place diversity')
#plt.ylabel('Friendship probability')
#plt.savefig('place_div_correlation.pdf')
#plt.close()

plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
lbl = []
for month,m in zip(sorted(plots),markers):
    x1,x2,y1,y2,x3 = plots[month]
    a1,b1,s1 = bb(y2,x2,0.001,10)

#    c = scipy.corrcoef(x2,y2)
#    p = scipy.stats.pearsonr(x2,y2)
#    s = scipy.stats.spearmanr(x2,y2)
#    print c,p,s
#    c = c[0,1]
#    p = p[0]
#    s = s[0]
#    print 'month %d, correlation %f, pearson %f, spearman %f'%(
#            month,c,p,s)
    plt.semilogy(a1,b1,'k%s-'%m)
    lbl.append('Snapshot %d'%(month-4))
plt.legend(lbl,loc='lower left',numpoints=1)
plt.grid(True)
plt.axis((0,9,1e-4,1))
plt.xlabel('Place entropy [bits]')
plt.ylabel('Link probability')
plt.savefig('place_ent_correlation.pdf')
plt.close()

plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
lbl = []
for month,m in zip(sorted(plots),markers):
    x1,x2,y1,y2,x3 = plots[month]
    a1,b1,s1 = bb(x3,x2,1,1e5,log=True)
    plt.loglog(a1,b1,'k%s-'%m)
    lbl.append('Snapshot %d'%(month-4))
plt.legend(lbl,loc='lower left',numpoints=1)
plt.grid(True)
#plt.xlim(0,9)
plt.axis((1,1e4,1e-4,1))
plt.xlabel('Number of check-ins')
plt.ylabel('Link probability')
plt.savefig('place_chk_correlation.pdf')
plt.close()

plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
lbl = []
for month,m in zip(sorted(plots),markers):
    x1,x2,y1,y2,x3 = plots[month]
    a1,b1,s1 = bb(x1,x2,1,1e4,log=True)
    plt.loglog(a1,b1,'k%s-'%m)
    lbl.append('Snapshot %d'%(month-4))
plt.legend(lbl,loc='lower left',numpoints=1)
plt.grid(True)
plt.axis((1,1e4,1e-4,1))
#plt.xlim(0,9)
plt.xlabel('Number of users')
plt.ylabel('Link probability')
plt.savefig('place_usr_correlation.pdf')
plt.close()
